Skip to content

[DRAFT] Add basic test for GQA (fusion)#2142

Closed
gramalingam wants to merge 11 commits intomainfrom
rama/gqa-basic-test
Closed

[DRAFT] Add basic test for GQA (fusion)#2142
gramalingam wants to merge 11 commits intomainfrom
rama/gqa-basic-test

Conversation

@gramalingam
Copy link
Collaborator

@gramalingam gramalingam commented Mar 27, 2025

Adds a test to verify equivalence of an expanded graph and GQA operator (which will serve as the basis for a fusion rule).

@codecov
Copy link

codecov bot commented Mar 27, 2025 •
edited
Loading

❌ 7 Tests Failed:

Tests completed Failed Passed Skipped
13268 7 13261 2405
View the top 3 failed test(s) by shortest run time 
onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0509_test_layer_normalization_2d_axis_negative_1
Stack Traces | 0.004s run time 
onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
C:\hostedtoolcache\windows\Python\3.11.9\x64\Lib\importlib\__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_layer_normalization_2d_axis_negative_1'

The above exception was the direct cause of the following exception:
.nox\test_onnx_weekly\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_layer_normalization_2d_axis_negative_1' (e=No module named 'tests.onnx_backend_test_code.test_layer_normalization_2d_axis_negative_1') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_layer_normalization_2d_axis_negative_1.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_layer_normalization_2d_axis_negative_1.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import FLOAT
E   from onnxscript.onnx_opset import opset17
E   
E   @script()
E   def bck_test_layer_normalization_2d_axis_negative_1(X: FLOAT[3,4], W: FLOAT[4], B: FLOAT[4]) -> (FLOAT[3,4], FLOAT[3,1], FLOAT[3,1]):
E       Y, Mean, InvStdDev = opset17.LayerNormalization(X, W, B, axis=-1)
E       return Y, Mean, InvStdDev
onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0844_test_reduce_log_sum_empty_set_expanded
Stack Traces | 0.004s run time 
onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
C:\hostedtoolcache\windows\Python\3.12.9\x64\Lib\importlib\__init__.py:90: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_reduce_log_sum_empty_set_expanded'

The above exception was the direct cause of the following exception:
.nox\test\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_reduce_log_sum_empty_set_expanded' (e=No module named 'tests.onnx_backend_test_code.test_reduce_log_sum_empty_set_expanded') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_reduce_log_sum_empty_set_expanded.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_reduce_log_sum_empty_set_expanded.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import FLOAT, INT64
E   from onnxscript.onnx_opset import opset18
E   
E   @script()
E   def bck_test_reduce_log_sum_empty_set_expanded(data: FLOAT[2,0,4], axes: INT64[1]) -> (FLOAT[2,1,4]):
E       ReduceLogSum_test_reduce_log_sum_empty_set_expanded_function_reduced_sum = opset18.ReduceSum(data, axes, keepdims=1)
E       reduced = opset18.Log(ReduceLogSum_test_reduce_log_sum_empty_set_expanded_function_reduced_sum)
E       return reduced
onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0976_test_resize_tf_crop_and_resize_extrapolation_value
Stack Traces | 0.004s run time 
onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
C:\hostedtoolcache\windows\Python\3.11.9\x64\Lib\importlib\__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_resize_tf_crop_and_resize_extrapolation_value'

The above exception was the direct cause of the following exception:
.nox\test_torch_nightly\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_resize_tf_crop_and_resize_extrapolation_value' (e=No module named 'tests.onnx_backend_test_code.test_resize_tf_crop_and_resize_extrapolation_value') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_resize_tf_crop_and_resize_extrapolation_value.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_resize_tf_crop_and_resize_extrapolation_value.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import FLOAT, INT64
E   from onnxscript.onnx_opset import opset19
E   
E   @script()
E   def bck_test_resize_tf_crop_and_resize_extrapolation_value(X: FLOAT[1,1,4,4], roi: FLOAT[8], sizes: INT64[4]) -> (FLOAT[1,1,3,3]):
E       Y = opset19.Resize(X, roi, None, sizes, coordinate_transformation_mode='tf_crop_and_resize', extrapolation_value=10.0, mode='linear')
E       return Y

To view more test analytics, go to the Test Analytics Dashboard
📋 Got 3 mins? Take this short survey to help us improve Test Analytics.

shubhambhokare1
shubhambhokare1 reviewed Mar 28, 2025
View reviewed changes
Copy link
Contributor

@shubhambhokare1 shubhambhokare1 left a comment •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Given it is a just a check to see whether the GQA operates in the same way as we think and not really checking the rewriting, maybe we can create a folder called something like op_equivalence_tests and place this there?

@gramalingam
Copy link
Collaborator Author

gramalingam commented Mar 29, 2025

nit: Given it is a just a check to see whether the GQA operates in the same way as we think and not really checking the rewriting, maybe we can create a folder called something like op_equivalence_tests and place this there?

Before discussing where to put this, let me explain why I created this. I have a GQA fusion implementation, which successfully replaces a subgraph with GQA ... but it produces a different output (with and without the fusion). It is very complicated pattern, and it was not clear what exactly was wrong. An advantage of a test-case like this one is that I can execute through the "expanded subgraph" and the 'replacement subgraph' using a debugger, stopping at intermediate points, to verify values look the way I think they should look (for example: does the attention mask look like a causal mask, etc.). (All I have to do is call the script-function with the inputs, instead of using its model-proto for an ORT session.)

The proof of correctness of a fusion rule comes from a proof of equivalence of code-fragments as captured in this test-case.

@gramalingam gramalingam mentioned this pull request Mar 30, 2025
Merged
@gramalingam
Copy link
Collaborator Author

gramalingam commented Mar 31, 2025

nit: Given it is a just a check to see whether the GQA operates in the same way as we think and not really checking the rewriting, maybe we can create a folder called something like op_equivalence_tests and place this there?

Before discussing where to put this, let me explain why I created this. I have a GQA fusion implementation, which successfully replaces a subgraph with GQA ... but it produces a different output (with and without the fusion). It is very complicated pattern, and it was not clear what exactly was wrong. An advantage of a test-case like this one is that I can execute through the "expanded subgraph" and the 'replacement subgraph' using a debugger, stopping at intermediate points, to verify values look the way I think they should look (for example: does the attention mask look like a causal mask, etc.). (All I have to do is call the script-function with the inputs, instead of using its model-proto for an ORT session.)

The proof of correctness of a fusion rule comes from a proof of equivalence of code-fragments as captured in this test-case.

To continue the discussion: once a working fusion rule is in place, this test can be easily turned into a fusion rule (by applying the fusion to the expanded model defined in the test case).

@gramalingam gramalingam changed the title Add basic test for GQA (fusion) [DRAFT] Add basic test for GQA (fusion) Apr 2, 2025
@gramalingam gramalingam marked this pull request as draft April 2, 2025 04:58
@gramalingam
Copy link
Collaborator Author

gramalingam commented Apr 2, 2025

Closing this. Will create an updated PR with GQA fusion added.

@gramalingam gramalingam closed this Apr 2, 2025
@github-project-automation github-project-automation bot moved this from Todo to Done in ONNX Script Review Board Apr 2, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

1 more reviewer

@shubhambhokare1 shubhambhokare1 shubhambhokare1 left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

None yet

Projects

ONNX Script Review Board
Status: Done

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@gramalingam @shubhambhokare1